package com.virtualperson.util;

import com.virtualperson.dto.ImageGenerationResponse;
import com.virtualperson.exception.ImageGenerationException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.concurrent.TimeUnit;

/**
 * 任务轮询工具类
 * 用于轮询任务状态，支持以下状态的轮询：
 * - PENDING（排队中）：任务在队列中等待处理
 * - RUNNING（处理中）：任务正在执行
 * - SUSPENDED（挂起）：任务暂时挂起，等待继续处理
 * 
 * 轮询将持续进行，直到任务达到以下最终状态：
 * - SUCCESS（执行成功）：任务成功完成
 * - FAILED（执行失败）：任务执行失败
 * - UNKNOWN（状态未知）：任务状态无法确定
 */
@Slf4j
@Component
public class TaskPollingUtil {
    // 轮询配置常量
    private static final long DEFAULT_INITIAL_DELAY = 1;  // 初始延迟时间（秒）
    private static final long DEFAULT_MAX_DELAY = 30;     // 最大延迟时间（秒）
    private static final long DEFAULT_TIMEOUT = 300;      // 超时时间（秒）
    private static final double BACKOFF_MULTIPLIER = 1.5; // 退避乘数

    // 任务状态常量
    private static final String STATUS_SUCCESS = "SUCCESS";
    private static final String STATUS_FAILED = "FAILED";
    private static final String STATUS_SUSPENDED = "SUSPENDED";
    private static final String STATUS_PENDING = "PENDING";
    private static final String STATUS_RUNNING = "RUNNING";
    private static final String STATUS_UNKNOWN = "UNKNOWN";

    /**
     * 轮询任务状态直到完成或超时
     * 对于PENDING（排队中）、RUNNING（处理中）、SUSPENDED（挂起）状态的任务会持续轮询
     * 
     * @param taskId 任务ID
     * @param taskQueryFunction 任务查询函数
     * @param statusCheck 状态检查函数
     * @param maxAttempts 最大轮询次数
     * @param intervalSeconds 轮询间隔（秒）
     * @return 最终的任务响应
     * @throws ImageGenerationException 如果任务失败或超时
     */
    public static ImageGenerationResponse pollTaskStatus(
            String taskId,
            TaskQueryFunction taskQueryFunction,
            java.util.function.Predicate<String> statusCheck,
            int maxAttempts,
            int intervalSeconds) {
        
        int attempts = 0;
        while (attempts < maxAttempts) {
            try {
                // 查询任务状态
                ImageGenerationResponse response = taskQueryFunction.queryTask(taskId);
                String status = response.getStatus();
                
                // 检查是否达到终止状态
                if (statusCheck.test(status)) {
                    return response;
                }
                
                // 根据状态处理
                switch (status) {
                    case STATUS_SUCCESS:
                        log.info("任务[{}]执行成功，轮询结束", taskId);
                        return response;
                    case STATUS_FAILED:
                        String errorMsg = String.format("任务[%s]执行失败，轮询结束", taskId);
                        log.error(errorMsg);
                        throw new ImageGenerationException(errorMsg);
                    case STATUS_UNKNOWN:
                        String unknownMsg = String.format("任务[%s]状态未知，轮询结束", taskId);
                        log.error(unknownMsg);
                        throw new ImageGenerationException(unknownMsg);
                    case STATUS_PENDING:
                        log.info("任务[{}]正在排队中，{}秒后重试", taskId, intervalSeconds);
                        break;
                    case STATUS_RUNNING:
                        log.info("任务[{}]正在处理中，{}秒后重试", taskId, intervalSeconds);
                        break;
                    case STATUS_SUSPENDED:
                        log.info("任务[{}]已挂起，{}秒后重试", taskId, intervalSeconds);
                        break;
                    default:
                        log.warn("任务[{}]状态未知: {}，{}秒后重试", taskId, status, intervalSeconds);
                }
                
                // 等待指定的间隔时间
                Thread.sleep(TimeUnit.SECONDS.toMillis(intervalSeconds));
                attempts++;
                
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                String errorMsg = String.format("任务[%s]轮询被中断", taskId);
                log.error(errorMsg, e);
                throw new ImageGenerationException(errorMsg);
            } catch (Exception e) {
                String errorMsg = String.format("任务[%s]轮询发生异常", taskId);
                log.error(errorMsg, e);
                throw new ImageGenerationException(errorMsg, e);
            }
        }
        
        // 如果达到最大尝试次数仍未完成
        String timeoutMsg = String.format("任务[%s]轮询超时，已尝试%d次", taskId, maxAttempts);
        log.error(timeoutMsg);
        throw new ImageGenerationException(timeoutMsg);
    }

    /**
     * 任务查询函数接口
     */
    @FunctionalInterface
    public interface TaskQueryFunction {
        /**
         * 查询任务状态
         * @param taskId 任务ID
         * @return 任务响应信息
         * @throws Exception 查询过程中可能发生的异常
         */
        ImageGenerationResponse queryTask(String taskId) throws Exception;
    }
} 